module scrypt

import crypto.pbkdf2
import crypto.sha256

fn test_salsa20_8() {
	// The input_block and output_block values are taken from
	// [RFC7914](https://datatracker.ietf.org/doc/html/rfc7914#section-8)
	// section 8.

	// vfmt off
    mut input_block := [
        u8(0x7e), 0x87, 0x9a, 0x21,   0x4f, 0x3e, 0xc9, 0x86,
           0x7c,  0xa9, 0x40, 0xe6,   0x41, 0x71, 0x8f, 0x26,
           0xba,  0xee, 0x55, 0x5b,   0x8c, 0x61, 0xc1, 0xb5,
           0x0d,  0xf8, 0x46, 0x11,   0x6d, 0xcd, 0x3b, 0x1d,
           0xee,  0x24, 0xf3, 0x19,   0xdf, 0x9b, 0x3d, 0x85,
           0x14,  0x12, 0x1e, 0x4b,   0x5a, 0xc5, 0xaa, 0x32,
           0x76,  0x02, 0x1d, 0x29,   0x09, 0xc7, 0x48, 0x29,
           0xed,  0xeb, 0xc6, 0x8d,   0xb8, 0xb8, 0xc2, 0x5e]

    output_block := [
        u8(0xa4), 0x1f, 0x85, 0x9c,   0x66, 0x08, 0xcc, 0x99,
           0x3b,  0x81, 0xca, 0xcb,   0x02, 0x0c, 0xef, 0x05,
           0x04,  0x4b, 0x21, 0x81,   0xa2, 0xfd, 0x33, 0x7d,
           0xfd,  0x7b, 0x1c, 0x63,   0x96, 0x68, 0x2f, 0x29,
           0xb4,  0x39, 0x31, 0x68,   0xe3, 0xc9, 0xe6, 0xbc,
           0xfe,  0x6b, 0xc5, 0xb7,   0xa0, 0x6d, 0x96, 0xba,
           0xe4,  0x24, 0xcc, 0x10,   0x2c, 0x91, 0x74, 0x5c,
           0x24,  0xad, 0x67, 0x3d,   0xc7, 0x61, 0x8f, 0x81]
	// vfmt on

	salsa20_8(mut input_block)

	for i in 0 .. 64 {
		assert input_block[i] == output_block[i], 'assertion failed for i: ${i}'
	}
}

fn test_block_mix() {
	// The input_block and output_block values are taken from
	// [RFC7914](https://datatracker.ietf.org/doc/html/rfc7914#section-9)
	// section 9.

	// vfmt off
    mut input_block := [
        // B[0] - the first 64 bytes of input
        u8(0xf7), 0xce, 0x0b, 0x65,   0x3d, 0x2d, 0x72, 0xa4,
           0x10,  0x8c, 0xf5, 0xab,   0xe9, 0x12, 0xff, 0xdd,
           0x77,  0x76, 0x16, 0xdb,   0xbb, 0x27, 0xa7, 0x0e,
           0x82,  0x04, 0xf3, 0xae,   0x2d, 0x0f, 0x6f, 0xad,
           0x89,  0xf6, 0x8f, 0x48,   0x11, 0xd1, 0xe8, 0x7b,
           0xcc,  0x3b, 0xd7, 0x40,   0x0a, 0x9f, 0xfd, 0x29,
           0x09,  0x4f, 0x01, 0x84,   0x63, 0x95, 0x74, 0xf3,
           0x9a,  0xe5, 0xa1, 0x31,   0x52, 0x17, 0xbc, 0xd7,
        // B[1] - the second 64 bytes of input
           0x89,  0x49, 0x91, 0x44,   0x72, 0x13, 0xbb, 0x22,
           0x6c,  0x25, 0xb5, 0x4d,   0xa8, 0x63, 0x70, 0xfb,
           0xcd,  0x98, 0x43, 0x80,   0x37, 0x46, 0x66, 0xbb,
           0x8f,  0xfc, 0xb5, 0xbf,   0x40, 0xc2, 0x54, 0xb0,
           0x67,  0xd2, 0x7c, 0x51,   0xce, 0x4a, 0xd5, 0xfe,
           0xd8,  0x29, 0xc9, 0x0b,   0x50, 0x5a, 0x57, 0x1b,
           0x7f,  0x4d, 0x1c, 0xad,   0x6a, 0x52, 0x3c, 0xda,
           0x77,  0x0e, 0x67, 0xbc,   0xea, 0xaf, 0x7e, 0x89]

    output_block := [
        // B'[0] - the first 64 bytes of output
        u8(0xa4), 0x1f, 0x85, 0x9c,   0x66, 0x08, 0xcc, 0x99,
           0x3b,  0x81, 0xca, 0xcb,   0x02, 0x0c, 0xef, 0x05,
           0x04,  0x4b, 0x21, 0x81,   0xa2, 0xfd, 0x33, 0x7d,
           0xfd,  0x7b, 0x1c, 0x63,   0x96, 0x68, 0x2f, 0x29,
           0xb4,  0x39, 0x31, 0x68,   0xe3, 0xc9, 0xe6, 0xbc,
           0xfe,  0x6b, 0xc5, 0xb7,   0xa0, 0x6d, 0x96, 0xba,
           0xe4,  0x24, 0xcc, 0x10,   0x2c, 0x91, 0x74, 0x5c,
           0x24,  0xad, 0x67, 0x3d,   0xc7, 0x61, 0x8f, 0x81,
        // B'[1] - the second 64 bytes of output
           0x20,  0xed, 0xc9, 0x75,   0x32, 0x38, 0x81, 0xa8,
           0x05,  0x40, 0xf6, 0x4c,   0x16, 0x2d, 0xcd, 0x3c,
           0x21,  0x07, 0x7c, 0xfe,   0x5f, 0x8d, 0x5f, 0xe2,
           0xb1,  0xa4, 0x16, 0x8f,   0x95, 0x36, 0x78, 0xb7,
           0x7d,  0x3b, 0x3d, 0x80,   0x3b, 0x60, 0xe4, 0xab,
           0x92,  0x09, 0x96, 0xe5,   0x9b, 0x4d, 0x53, 0xb6,
           0x5d,  0x2a, 0x22, 0x58,   0x77, 0xd5, 0xed, 0xf5,
           0x84,  0x2c, 0xb9, 0xf1,   0x4e, 0xef, 0xe4, 0x25]
	// vfmt on

	// an array capable of holding r * 128 bytes used during
	// the block_mix operation.
	mut temp_block := []u8{len: 128, cap: 128}

	// for this test, r = 1
	block_mix(mut input_block, mut temp_block, 1)

	for i in 0 .. 128 {
		assert input_block[i] == output_block[i], 'assertion failed for i: ${i}'
	}
}

fn test_smix() {
	// The input_block and output_block values are taken from
	// [RFC7914](https://datatracker.ietf.org/doc/html/rfc7914#section-10)
	// section 10.

	// vfmt off
    mut input_block := [
        u8(0xf7), 0xce, 0x0b, 0x65,   0x3d, 0x2d, 0x72, 0xa4,
           0x10,  0x8c, 0xf5, 0xab,   0xe9, 0x12, 0xff, 0xdd,
           0x77,  0x76, 0x16, 0xdb,   0xbb, 0x27, 0xa7, 0x0e,
           0x82,  0x04, 0xf3, 0xae,   0x2d, 0x0f, 0x6f, 0xad,
           0x89,  0xf6, 0x8f, 0x48,   0x11, 0xd1, 0xe8, 0x7b,
           0xcc,  0x3b, 0xd7, 0x40,   0x0a, 0x9f, 0xfd, 0x29,
           0x09,  0x4f, 0x01, 0x84,   0x63, 0x95, 0x74, 0xf3,
           0x9a,  0xe5, 0xa1, 0x31,   0x52, 0x17, 0xbc, 0xd7,
           0x89,  0x49, 0x91, 0x44,   0x72, 0x13, 0xbb, 0x22,
           0x6c,  0x25, 0xb5, 0x4d,   0xa8, 0x63, 0x70, 0xfb,
           0xcd,  0x98, 0x43, 0x80,   0x37, 0x46, 0x66, 0xbb,
           0x8f,  0xfc, 0xb5, 0xbf,   0x40, 0xc2, 0x54, 0xb0,
           0x67,  0xd2, 0x7c, 0x51,   0xce, 0x4a, 0xd5, 0xfe,
           0xd8,  0x29, 0xc9, 0x0b,   0x50, 0x5a, 0x57, 0x1b,
           0x7f,  0x4d, 0x1c, 0xad,   0x6a, 0x52, 0x3c, 0xda,
           0x77,  0x0e, 0x67, 0xbc,   0xea, 0xaf, 0x7e, 0x89]

    output_block := [
        u8(0x79), 0xcc, 0xc1, 0x93,   0x62, 0x9d, 0xeb, 0xca,
           0x04,  0x7f, 0x0b, 0x70,   0x60, 0x4b, 0xf6, 0xb6,
           0x2c,  0xe3, 0xdd, 0x4a,   0x96, 0x26, 0xe3, 0x55,
           0xfa,  0xfc, 0x61, 0x98,   0xe6, 0xea, 0x2b, 0x46,
           0xd5,  0x84, 0x13, 0x67,   0x3b, 0x99, 0xb0, 0x29,
           0xd6,  0x65, 0xc3, 0x57,   0x60, 0x1f, 0xb4, 0x26,
           0xa0,  0xb2, 0xf4, 0xbb,   0xa2, 0x00, 0xee, 0x9f,
           0x0a,  0x43, 0xd1, 0x9b,   0x57, 0x1a, 0x9c, 0x71,
           0xef,  0x11, 0x42, 0xe6,   0x5d, 0x5a, 0x26, 0x6f,
           0xdd,  0xca, 0x83, 0x2c,   0xe5, 0x9f, 0xaa, 0x7c,
           0xac,  0x0b, 0x9c, 0xf1,   0xbe, 0x2b, 0xff, 0xca,
           0x30,  0x0d, 0x01, 0xee,   0x38, 0x76, 0x19, 0xc4,
           0xae,  0x12, 0xfd, 0x44,   0x38, 0xf2, 0x03, 0xa0,
           0xe4,  0xe1, 0xc4, 0x7e,   0xc3, 0x14, 0x86, 0x1f,
           0x4e,  0x90, 0x87, 0xcb,   0x33, 0x39, 0x6a, 0x68,
           0x73,  0xe8, 0xf9, 0xd2,   0x53, 0x9a, 0x4b, 0x8e]
	// vfmt on

	r := u32(1)
	n := u64(16)

	// len and cap are 128 * r * n = 2048
	mut v_block := []u8{len: 2048, cap: 2048}

	// len and cap are 256 * r = 246
	mut temp_block := []u8{len: 256, cap: 256}

	smix(mut input_block, r, n, mut v_block, mut temp_block)

	for i in 0 .. 128 {
		assert input_block[i] == output_block[i], 'assertion failed for i: ${i}'
	}
}

fn test_pbkdf2_hmac_sha256() {
	// The input_block and output_block values are taken from
	// [RFC7914](https://datatracker.ietf.org/doc/html/rfc7914#section-11)
	// section 11.

	// vfmt off
    output_block := [
        [u8(0x55), 0xac, 0x04, 0x6e,   0x56, 0xe3, 0x08, 0x9f,
            0xec,  0x16, 0x91, 0xc2,   0x25, 0x44, 0xb6, 0x05,
            0xf9,  0x41, 0x85, 0x21,   0x6d, 0xde, 0x04, 0x65,
            0xe6,  0x8b, 0x9d, 0x57,   0xc2, 0x0d, 0xac, 0xbc,
            0x49,  0xca, 0x9c, 0xcc,   0xf1, 0x79, 0xb6, 0x45,
            0x99,  0x16, 0x64, 0xb3,   0x9d, 0x77, 0xef, 0x31,
            0x7c,  0x71, 0xb8, 0x45,   0xb1, 0xe3, 0x0b, 0xd5,
            0x09,  0x11, 0x20, 0x41,   0xd3, 0xa1, 0x97, 0x83
        ],
        [u8(0x4d), 0xdc, 0xd8, 0xf6,   0x0b, 0x98, 0xbe, 0x21,
            0x83,  0x0c, 0xee, 0x5e,   0xf2, 0x27, 0x01, 0xf9,
            0x64,  0x1a, 0x44, 0x18,   0xd0, 0x4c, 0x04, 0x14,
            0xae,  0xff, 0x08, 0x87,   0x6b, 0x34, 0xab, 0x56,
            0xa1,  0xd4, 0x25, 0xa1,   0x22, 0x58, 0x33, 0x54,
            0x9a,  0xdb, 0x84, 0x1b,   0x51, 0xc9, 0xb3, 0x17,
            0x6a,  0x27, 0x2b, 0xde,   0xbb, 0xa1, 0xd0, 0x78,
            0x47,  0x8f, 0x62, 0xb3,   0x97, 0xf3, 0x3c, 0x8d
        ]
    ]
	// vfmt on

	d0 := pbkdf2.key('passwd'.bytes(), 'salt'.bytes(), 1, 64, sha256.new())!

	assert d0 == output_block[0]

	d1 := pbkdf2.key('Password'.bytes(), 'NaCl'.bytes(), 80000, 64, sha256.new())!

	assert d1 == output_block[1]
}

struct ScryptTestData {
	name            string
	password        []u8
	salt            []u8
	n               u64
	r               u32
	p               u32
	dk_len          u64
	expected_result []u8
}

// The scrypt test vectors are taken from
// [RFC7914](https://datatracker.ietf.org/doc/html/rfc7914#section-12)
// section 12.
const scrypt_test_cases = [
	ScryptTestData{
		name:            'test case 1'
		password:        ''.bytes()
		salt:            ''.bytes()
		n:               16
		r:               1
		p:               1
		dk_len:          64
		expected_result: [u8(0x77), 0xd6, 0x57, 0x62, 0x38, 0x65, 0x7b, 0x20, 0x3b, 0x19, 0xca,
			0x42, 0xc1, 0x8a, 0x04, 0x97, 0xf1, 0x6b, 0x48, 0x44, 0xe3, 0x07, 0x4a, 0xe8, 0xdf,
			0xdf, 0xfa, 0x3f, 0xed, 0xe2, 0x14, 0x42, 0xfc, 0xd0, 0x06, 0x9d, 0xed, 0x09, 0x48,
			0xf8, 0x32, 0x6a, 0x75, 0x3a, 0x0f, 0xc8, 0x1f, 0x17, 0xe8, 0xd3, 0xe0, 0xfb, 0x2e,
			0x0d, 0x36, 0x28, 0xcf, 0x35, 0xe2, 0x0c, 0x38, 0xd1, 0x89, 0x06]
	},
	// test cases 2, 3, and 4 are moved to the slow test repo.
]

fn test_scrypt() {
	for c in scrypt_test_cases {
		results := scrypt(c.password, c.salt, c.n, c.r, c.p, c.dk_len)!
		assert results == c.expected_result, '${c.name} failed'
	}
}
